# %%
import SPTAG
import shutil
import numpy as np
import struct
import os

# %%定义运行参数
query_path = "/home/ljl/Code/dataset/vector-ssd/bigann/bigann_query.bbin"
dataset_type_name = "Int8" #Float, Int8
index_path = "sptag_index"
dataset_type = np.float32 if dataset_type_name == "Float" else np.uint8
threads = 32

# %%
def read_bin(file_path, type):
    file_size = os.path.getsize(file_path)
    print("Read data from", file_path, "\nfile size:",file_size)
    feature_size = np.dtype(type).itemsize
    with open(file_path,"rb") as fd:
        lines = int.from_bytes(fd.read(4), byteorder='little')
        dim = int.from_bytes(fd.read(4), byteorder='little')
        print("lines:",lines,"dim:",dim)

        data_size = lines * dim * feature_size
        if(data_size+8 != file_size):
            raise Exception(f"Error! file size {file_size} and argument {data_size+8} not match!") # 判断实际文件大小是否与参数匹配，简单的纠错机制

        binary_data = fd.read(data_size)
        vectors = np.frombuffer(binary_data, dtype=type)
        vectors = vectors.reshape(lines, dim)
        print("Returned vector list:",vectors.shape, vectors.dtype)

        return vectors
    
# %% 读取数据集
querys = read_bin(query_path, dataset_type)
n = querys.shape[0]
dim = querys.shape[1]
print("querys shape:", n, dim)

# %%
index = SPTAG.AnnIndex.Load(index_path)

# %%
result = index.SearchWithMetaData(querys, 10)
# %%
print (result[0]) # nearest k vector ids
print (result[1]) # nearest k vector distances
print (result[2]) # nearest k vector metadatas
# %%
